from pathlib import Path
PATH_ROOT = Path(__file__).parent
import sys
sys.path.append(str(PATH_ROOT))
PATH_DATA = PATH_ROOT / "data"
PATH_DATA.mkdir(exist_ok=True)

import cv2
import time
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
from infer import InferSession
from lightglue_dynamo import viz

def noclip_rotation(image: np.ndarray, angle: float):
    """Rotate an image without clipping.

    Args:
        image (np.ndarray): shape=(H, W, C)
        angle (float): Rotation angle in degrees.
    Returns:
        Tuple[np.ndarray, np.ndarray, Tuple[int, int]]: Rotated image, rotation matrix, new shape (width, height).
    """
    h, w = image.shape[:2]
    center = (w // 2, h // 2)
    rot_mat = cv2.getRotationMatrix2D(center, angle, 1.0)
    cos = np.abs(rot_mat[0, 0])
    sin = np.abs(rot_mat[0, 1])
    
    new_w = int((h * sin) + (w * cos))
    new_h = int((h * cos) + (w * sin))
    
    rot_mat[0, 2] += (new_w / 2) - center[0]
    rot_mat[1, 2] += (new_h / 2) - center[1]
    
    rotated = cv2.warpAffine(image, rot_mat, (new_w, new_h))
    return rotated, rot_mat, (new_w, new_h)

def plot_images(imgs, titles=None, show=False, save_path: Path = None):
    n = len(imgs)
    fig, ax = plt.subplots(1, n, figsize=(n*5, 4))
    for i in range(n):
        ax[i].imshow(imgs[i])
        ax[i].set_title(titles[i] if titles else f"Image {i+1}")
        ax[i].axis("off")
    fig.tight_layout()
    fig.canvas.draw()
    w, h = fig.canvas.get_width_height()
    img_array = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8).reshape(h, w, 4)[:, :, 1:]
    if save_path is not None:
        plt.savefig(save_path, dpi=100)
    if show:
        plt.show()
    plt.close()
    return img_array

class MapAligner:
    def __init__(self, device='cpu'):
        self.infer_session = InferSession(device=device)
    
    def rotate_compare(self, image0: np.ndarray, image1: np.ndarray, delta_angle: float = 36, verbose: bool = False, save: bool = False):
        """Rotate and compare two images.

        Args:
            image0 (np.ndarray): shape=(H, W, C)
            image1 (np.ndarray): shape=(H, W, C)
            delta_angle (float, optional): Delta angle for rotation. Defaults to 36.
            verbose (bool, optional): If True, print detailed information. Defaults to False.
        
        Returns:
            Tuple[np.ndarray, int]: Affine transformation matrix and the number of matches.
        """
        assert image0.ndim == 3 and image1.ndim == 3, "Both images must be 3-dimensional (H, W, C)."
        angles = np.arange(0, 360, delta_angle)
        best = {
            'score': 0,
            'angle': None,
            'matches': None,
            'rot_mat': None,
            'shape': None
        }
        
        for angle in angles:
            rotated_image0, rot_mat, new_shape = noclip_rotation(image0, angle)
            # cv2.imshow(f"Rotated Image at {angle} degrees", rotated_image0)
            # cv2.waitKey(0)
            m_kpts0, m_kpts1 = self.infer_session(rotated_image0, image1)
            score = len(m_kpts0)
            if score > best['score']:
                best['score'] = score
                best['angle'] = angle
                best['matches'] = (m_kpts0, m_kpts1)
                best['rot_mat'] = rot_mat
                best['shape'] = new_shape
            if verbose:
                print(f"Angle: {angle} degrees, Matches: {score}")
        
        if verbose:
            print(f"Best angle: {best['angle']} degrees with {best['score']} matches")
        if best['matches'] is None:
            print("No matches found.")
            return None, 0, None

        A = cv2.estimateAffinePartial2D(best['matches'][0], best['matches'][1], method=cv2.RANSAC)[0]
        rot_mat = np.vstack([best['rot_mat'], [0, 0, 1]])
        A = A @ rot_mat
        A = A[:2, :]
        angle_rad = np.arctan2(A[1, 0], A[0, 0])
        translation = A[:, 2]

        if verbose:
            print(f"Rotation angle (rad): {angle_rad}, Translation: {translation}")
        if verbose or save:
            com_img = plot_images(
                [image0, cv2.warpAffine(image0, A, (image1.shape[1], image1.shape[0])), image1],
                [f"Original Image", f"Aligned Image ({best['angle']}°)", "Reference Image"],
                show=verbose,
                save_path=str(PATH_DATA/"aligned_map.png") if save else None
            )
        return A, best['score'], com_img


if __name__ == '__main__':
    map_aligner = MapAligner('cuda')

    img2 = cv2.imread(str(PATH_ROOT/"assets/2d_map_-1868_-811.png"))
    img1 = cv2.imread(str(PATH_ROOT/"assets/2d_map_-1868_-811_3.png"))
    assert img1 is not None and img2 is not None, "Images not found. Please check the paths."

    start_time = time.time()
    A, score, com_img = map_aligner.rotate_compare(img1, img2, delta_angle=36, verbose=True)
    print(f"Time taken: {time.time() - start_time:.2f} seconds")
